
import cv2
import numpy as np
import onnxruntime as ort
import time
import socket
import cv2
import smbus
import math
import sys
from luma.core.interface.serial import i2c, spi
from luma.core.render import canvas
from luma.oled.device import ssd1306
from PIL import ImageFont

def plot_one_box(x, img, color=None, label=None, line_thickness=None):
    tl = (
        line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1
    )  # line/font thickness
    color = color or [random.randint(0, 255) for _ in range(3)]
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
    if label:
        tf = max(tl - 1, 1)  # font thickness
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled
        cv2.putText(
            img,
            label,
            (c1[0], c1[1] - 2),
            0,
            tl / 3,
            [225, 255, 255],
            thickness=tf,
            lineType=cv2.LINE_AA,
        )
 
def _make_grid( nx, ny):
        xv, yv = np.meshgrid(np.arange(ny), np.arange(nx))
        return np.stack((xv, yv), 2).reshape((-1, 2)).astype(np.float32)
 
def cal_outputs(outs,nl,na,model_w,model_h,anchor_grid,stride):
    
    row_ind = 0
    grid = [np.zeros(1)] * nl
    for i in range(nl):
        h, w = int(model_w/ stride[i]), int(model_h / stride[i])
        length = int(na * h * w)
        if grid[i].shape[2:4] != (h, w):
            grid[i] = _make_grid(w, h)
 
        outs[row_ind:row_ind + length, 0:2] = (outs[row_ind:row_ind + length, 0:2] * 2. - 0.5 + np.tile(
            grid[i], (na, 1))) * int(stride[i])
        outs[row_ind:row_ind + length, 2:4] = (outs[row_ind:row_ind + length, 2:4] * 2) ** 2 * np.repeat(
            anchor_grid[i], h * w, axis=0)
        row_ind += length
    return outs
 
 
 
def post_process_opencv(outputs,model_h,model_w,img_h,img_w,thred_nms,thred_cond):
    conf = outputs[:,4].tolist()
    c_x = outputs[:,0]/model_w*img_w
    c_y = outputs[:,1]/model_h*img_h
    w  = outputs[:,2]/model_w*img_w
    h  = outputs[:,3]/model_h*img_h
    p_cls = outputs[:,5:]
    if len(p_cls.shape)==1:
        p_cls = np.expand_dims(p_cls,1)
    cls_id = np.argmax(p_cls,axis=1)
 
    p_x1 = np.expand_dims(c_x-w/2,-1)
    p_y1 = np.expand_dims(c_y-h/2,-1)
    p_x2 = np.expand_dims(c_x+w/2,-1)
    p_y2 = np.expand_dims(c_y+h/2,-1)
    areas = np.concatenate((p_x1,p_y1,p_x2,p_y2),axis=-1)
    
    areas = areas.tolist()
    ids = cv2.dnn.NMSBoxes(areas,conf,thred_cond,thred_nms)
    if len(ids)>0:
        return  np.array(areas)[ids],np.array(conf)[ids],cls_id[ids]
    else:
        return [],[],[]
def infer_img(img0,net,model_h,model_w,nl,na,stride,anchor_grid,thred_nms=0.4,thred_cond=0.5):
    # 图像预处理
    img = cv2.resize(img0, [model_w,model_h], interpolation=cv2.INTER_AREA)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.astype(np.float32) / 255.0
    blob = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0)
 
    # 模型推理
    outs = net.run(None, {net.get_inputs()[0].name: blob})[0].squeeze(axis=0)
 
    # 输出坐标矫正
    outs = cal_outputs(outs,nl,na,model_w,model_h,anchor_grid,stride)
 
    # 检测框计算
    img_h,img_w,_ = np.shape(img0)
    boxes,confs,ids = post_process_opencv(outs,model_h,model_w,img_h,img_w,thred_nms,thred_cond)
 
    return  boxes,confs,ids
 
def SendVideo(sth):
    address = ('183.230.40.40', 1811) # 地址:IP+端口号
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) # 创建sock对象
        sock.connect(address) # 与客户端连接
    except socket.error as msg: # 判断连接是否发生错误
        print(msg)
        sys.exit(1)
    DataBase=b'*627214#spark123#sample*' #输入连接报文
    sock.send(DataBase)
    encoded_sth = sth.encode('utf-8')
    sock.send(encoded_sth) #发送数据
  
def oled_func(my_string:str,pos:str):
    serial = i2c(port=1, address=0x3C)
    device = ssd1306(serial)
    font = ImageFont.truetype('/home/derrickcr/simsun/simsun.ttc', 12)
    with canvas(device) as draw:
        draw.rectangle(device.bounding_box, outline="white", fill="black")
        draw.text((5, 10),my_string, fill="white", font=font)
        draw.text((10,10),pos,fill="white",font=font)

 # model_pb_path 和 dic_labels需要修改
 
if __name__ == "__main__":
    '是否发送信息/通过oled显示'
    bool_sendVideo = 1
    bool_oled_func = 0
    # 模型加载
    model_pb_path = "best_rubbish.onnx"
    so = ort.SessionOptions()
    net = ort.InferenceSession(model_pb_path, so)
    
    # 标签字典
    hazardous_waste_list = [1, 9, 10, 11, 14, 15, 22, 23, 27, 31, 36, 40, 41, 42, 46, 47, 51, 52, 56, 58]
    recyclable_waste_list = [2, 16, 17, 25, 26, 29, 32, 33, 35, 37, 43, 44, 48, 53, 54, 59]
    kitchen_waste_list = [3, 4, 5, 6, 7, 20, 21, 24, 28, 30, 38]
    other_waste_list = [0, 8, 13, 18, 19, 34, 39, 45, 49, 50, 55, 57]
    # 模型参数
    model_h = 320
    model_w = 320
    nl = 3
    na = 3
    stride=[8.,16.,32.]
    anchors = [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]]
    anchor_grid = np.asarray(anchors, dtype=np.float32).reshape(nl, -1, 2)
    
    video = 0
    cap = cv2.VideoCapture(video)
    flag_det = False
    while True:
        success, img0 = cap.read()
        if success:
            
            if flag_det:
                t1 = time.time()
                det_boxes,scores,ids = infer_img(img0,net,model_h,model_w,nl,na,stride,anchor_grid,thred_nms=0.4,thred_cond=0.5)
                for box,score,id in zip(det_boxes,scores,ids):
                    x1, y1, x2, y2 = box
                    width = x2 - x1
                    height = y2 - y1
                    c_x = (x1 + x2) / 2
                    c_y = (y1 + y2) / 2
                    position = '方框高为%.1f，宽为%.1f，中心点坐标为(%.1f,%.1f)'%(height, width, c_x, c_y)
                    print(position)
                    classify = ''
                    if id in hazardous_waste_list:
                        classify = '有害垃圾'
                        classify_e = 'hazardous_waste'
                    elif id in recyclable_waste_list:
                        classify = '可回收垃圾'
                        classify_e = 'recyclable_waste'
                    elif id in kitchen_waste_list:
                        classify = '厨余垃圾'
                        classify_e = 'kitchen_waste'
                    else:
                        classify = '其它垃圾'
                        classify_e = 'other_waste'
                    label_chinese = '%s:%.2f'%(classify,score)
                    label = '%s:%.2f'%(classify_e,score)
                    print(label)
                    if bool_sendVideo:
                        SendVideo(label_chinese)
                    if bool_oled_func:
                        oled_func(label_chinese,position)
                    plot_one_box(box.astype(np.int16), img0, color=(255,0,0), label=label, line_thickness=None)
                t2 = time.time()
                str_FPS = "FPS: %.2f"%(1./(t2-t1))
                cv2.putText(img0,str_FPS,(50,50),cv2.FONT_HERSHEY_COMPLEX,1,(0,255,0),3)
                
            
            cv2.imshow("video",img0)
        key=cv2.waitKey(1) & 0xFF    
        if key == ord('q'):
        
            break
        elif key & 0xFF == ord('s'):
            flag_det = not flag_det
            print(flag_det)
    if bool_oled_func:
        device.clear()
    cap.release()
